package io.github.hadyang.leetcode.offer;

import java.util.ArrayList;
import org.junit.Test;

/**
 * 输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4,。
 *
 * @author haoyang.shi
 */
public class GetLeastNumbers {

  @Test
  public void test() {
    System.out.println(GetLeastNumbers_Solution_Partition(new int[] {4, 5, 1, 6, 2, 7, 3, 8}, 1));
    System.out.println(GetLeastNumbers_Solution_Partition(new int[] {4, 5, 1, 6, 2, 7, 3, 8}, 8));
  }

  public ArrayList<Integer> GetLeastNumbers_Solution_Partition(int[] input, int k) {
    ArrayList<Integer> res = new ArrayList<>();

    if (k > input.length || k < 1) {
      return res;
    }

    int start = 0, end = input.length - 1;
    int index = partition(input, start, end);
    while (index != k - 1) {
      if (index > k - 1) {
        end = index - 1;
        index = partition(input, start, end);
      } else {
        start = index + 1;
        index = partition(input, start, end);
      }
    }

    for (int i = 0; i < input.length && i < k; i++) {
      res.add(input[i]);
    }
    return res;
  }

  private int partition(int[] nums, int start, int end) {
    int left = start, right = end;
    int key = nums[left];

    while (left < right) {
      while (left < right && nums[right] > key) {
        right--;
      }
      if (left < right) {
        nums[left] = nums[right];
        left++;
      }

      while (left < right && nums[left] <= key) {
        left++;
      }
      if (left < right) {
        nums[right] = nums[left];
        right++;
      }
    }

    nums[left] = key;

    return left;
  }

  public ArrayList<Integer> GetLeastNumbers_Solution_Heap(int[] input, int k) {
    ArrayList<Integer> res = new ArrayList<>();
    if (k > input.length || k == 0) {
      return res;
    }

    for (int i = input.length - 1; i >= 0; i--) {
      minHeap(input, 0, i);

      swap(input, 0, i);

      res.add(input[i]);
      if (res.size() == k) break;
    }
    return res;
  }

  private void minHeap(int[] heap, int start, int end) {
    if (start == end) {
      return;
    }

    int childLeft = start * 2 + 1;
    int childRight = childLeft + 1;

    if (childLeft <= end) {
      minHeap(heap, childLeft, end);

      if (heap[childLeft] < heap[start]) {
        swap(heap, start, childLeft);
      }
    }

    if (childRight <= end) {
      minHeap(heap, childRight, end);

      if (heap[childRight] < heap[start]) {
        swap(heap, start, childRight);
      }
    }
  }

  private void swap(int[] nums, int a, int b) {
    int t = nums[a];
    nums[a] = nums[b];
    nums[b] = t;
  }
}
